#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import rospy
import numpy as np
import tf.transformations as tft
from sensor_msgs.msg import Joy
from geometry_msgs.msg import PoseStamped
from trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint
from moveit_commander import MoveGroupCommander, roscpp_initialize, roscpp_shutdown
from moveit_msgs.srv import GetPositionIK, GetPositionIKRequest, GetPositionFK, GetPositionFKRequest
from copy import deepcopy


class TeleopPose:
    def __init__(self):
        roscpp_initialize([])
        rospy.init_node("teleop_pose_first_person_ik")

        # 参数加载
        self.group_name = rospy.get_param("~group_name", "robot_arm")
        self.step_xyz = rospy.get_param("~step_xyz", 0.002)  # 位置步长（米）
        self.step_rpy = rospy.get_param("~step_rpy", 0.5)    # 姿态步长（度）
        self.loop_rate = rospy.get_param("~loop_rate", 20)
        self.joy_deadzone = rospy.get_param("~joy_deadzone", 0.1)
        self.orientation_tol = rospy.get_param("~orientation_tolerance", 1.0)

        # MoveIt 初始化
        self.group = MoveGroupCommander(self.group_name)
        self.ee_link = self.group.get_end_effector_link()
        self.joint_names = self.group.get_joints()

        # 初始化目标位姿为当前位姿（以当前姿态为基准）
        current_pose = self.group.get_current_pose()
        self.target_pose = PoseStamped()
        self.target_pose.header = current_pose.header
        self.target_pose.pose = current_pose.pose

        self.print_pose(self.target_pose, "初始位姿（以当前姿态为基准）")

        # ROS 通信接口
        # 扩展为8个轴（索引0-7），匹配手柄实际轴数
        self.joy_axes = [0.0] * 8  # 0-7轴：根据你的定义映射
        self.joy_buttons = [0] * 12  # 按钮（用于姿态重置）
        rospy.Subscriber("/joy", Joy, self.cb_joy)
        self.joint_pub = rospy.Publisher("/robot_arm_control/command", JointTrajectory, queue_size=10)
        rospy.wait_for_service("/compute_ik")
        rospy.wait_for_service("/compute_fk")
        self.ik_srv = rospy.ServiceProxy("/compute_ik", GetPositionIK)
        self.fk_srv = rospy.ServiceProxy("/compute_fk", GetPositionFK)

        rospy.loginfo("✅ 遥操作已启动：末端平移+旋转（自定义轴映射）")
        rospy.loginfo("轴映射定义：")
        rospy.loginfo("  平移：请在update_target_position中定义（例如0:X, 1:Y, 3:Z）")
        rospy.loginfo("  旋转：axes[2]=Yaw, axes[6]=Pitch, axes[7]=Roll")
        self.run()

    # ========== 手柄回调 ==========
    def cb_joy(self, msg: Joy):
        # 限制轴索引在0-7范围内（防止手柄轴数不足导致错误）
        for i in range(min(8, len(msg.axes))):
            val = msg.axes[i]
            self.joy_axes[i] = val if abs(val) > self.joy_deadzone else 0.0
        
        # 保存按钮状态（按钮0用于姿态重置）
        self.joy_buttons = msg.buttons[:12]  # 限制为12个按钮

    # ========== 主循环 ==========
    def run(self):
        rate = rospy.Rate(self.loop_rate)
        while not rospy.is_shutdown():
            # 按钮0：重置目标姿态为当前实际姿态
            if self.joy_buttons[0] == 1:
                current_pose = self.group.get_current_pose()
                self.target_pose.pose = current_pose.pose
                rospy.loginfo("🔄 姿态已重置为当前实际姿态")

            if any(abs(v) > 1e-6 for v in self.joy_axes):
                prev_pose = deepcopy(self.target_pose.pose)  # 保存当前位姿用于回滚
                self.update_target_pose()                   # 更新目标位姿（位置+姿态）
                
                # 求解IK并发布，失败则回滚
                if not self.solve_and_publish():
                    self.target_pose.pose = prev_pose
                    rospy.logwarn("⚠️ IK失败 → 回滚位姿")
            rate.sleep()

    # ========== 位姿更新 ==========
    def update_target_pose(self):
        """更新目标位置（沿当前姿态坐标系）和姿态（按自定义轴映射）"""
        # 1. 位置更新（请根据你的手柄轴定义修改这里的索引）
        # 示例：假设axes[0]=X平移, axes[1]=Y平移, axes[3]=Z平移（可根据实际调整）
        dx = self.step_xyz * self.joy_axes[0]   # X方向平移
        dy = self.step_xyz * self.joy_axes[3]   # Y方向平移
        dz = self.step_xyz * self.joy_axes[1]   # Z方向平移

        # 沿当前姿态的坐标系计算位置增量
        current_rot = self.quat_to_rot(self.target_pose.pose.orientation)
        delta_pos = current_rot.dot(np.array([dx, dy, dz]))

        # 更新位置
        p = self.target_pose.pose.position
        p.x += delta_pos[0]
        p.y += delta_pos[1]
        p.z += delta_pos[2]

        # 2. 姿态更新（按你的需求：axes[2]=Yaw, axes[6]=Pitch, axes[7]=Roll）
        d_yaw = self.step_rpy * self.joy_axes[2]   # Yaw轴旋转（绕Z轴）
        d_pitch = self.step_rpy * self.joy_axes[6] # Pitch轴旋转（绕Y轴）
        d_roll = self.step_rpy * self.joy_axes[7]  # Roll轴旋转（绕X轴）

        # 将当前姿态转为欧拉角（弧度→度，方便计算）
        current_quat = [
            self.target_pose.pose.orientation.x,
            self.target_pose.pose.orientation.y,
            self.target_pose.pose.orientation.z,
            self.target_pose.pose.orientation.w
        ]
        current_rpy = np.rad2deg(tft.euler_from_quaternion(current_quat))  # [roll, pitch, yaw]

        # 应用旋转增量（注意：欧拉角顺序为roll→pitch→yaw）
        new_rpy = [
            current_rpy[0] + d_roll,   # Roll增量
            current_rpy[1] + d_pitch,  # Pitch增量
            current_rpy[2] + d_yaw     # Yaw增量
        ]

        # 转回四元数并更新目标姿态
        new_quat = tft.quaternion_from_euler(
            np.deg2rad(new_rpy[0]),  # roll（弧度）
            np.deg2rad(new_rpy[1]),  # pitch（弧度）
            np.deg2rad(new_rpy[2])   # yaw（弧度）
        )
        self.target_pose.pose.orientation.x = new_quat[0]
        self.target_pose.pose.orientation.y = new_quat[1]
        self.target_pose.pose.orientation.z = new_quat[2]
        self.target_pose.pose.orientation.w = new_quat[3]

        self.print_pose(self.target_pose, "目标位姿（更新后）", throttle=0.5)

    # ========== IK + FK 验证 ==========
    def solve_and_publish(self):
        req = GetPositionIKRequest()
        req.ik_request.group_name = self.group_name
        req.ik_request.pose_stamped = self.target_pose  # 使用当前目标位姿（含位置和姿态）
        req.ik_request.ik_link_name = self.ee_link
        req.ik_request.timeout = rospy.Duration(0.05)
        req.ik_request.robot_state.joint_state.name = self.joint_names
        req.ik_request.robot_state.joint_state.position = self.group.get_current_joint_values()

        try:
            res = self.ik_srv(req)
            if res.error_code.val != 1:
                rospy.logwarn_throttle(1.0, f"❌ IK失败 code={res.error_code.val}")
                return False
        except rospy.ServiceException as e:
            rospy.logerr(f"IK服务错误: {e}")
            return False

        # 验证FK计算的姿态与目标姿态偏差
        if not self.verify_orientation(res.solution.joint_state):
            rospy.logwarn("❌ 姿态偏差过大，未发布轨迹")
            return False

        self.publish_joint_traj(res.solution.joint_state)
        rospy.loginfo_throttle(0.5, "✅ 已发布IK解算结果")
        return True

    def verify_orientation(self, joint_state):
        fk_req = GetPositionFKRequest()
        fk_req.header.frame_id = self.target_pose.header.frame_id
        fk_req.fk_link_names = [self.ee_link]
        fk_req.robot_state.joint_state = joint_state
        try:
            fk_res = self.fk_srv(fk_req)
            if not fk_res.pose_stamped:
                rospy.logwarn("❌ FK失败")
                return False
            pose_calc = fk_res.pose_stamped[0].pose
        except rospy.ServiceException as e:
            rospy.logerr(f"FK服务错误: {e}")
            return False

        # 对比目标姿态与FK计算的姿态
        tgt_q, cal_q = self.target_pose.pose.orientation, pose_calc.orientation
        tgt_e = np.degrees(tft.euler_from_quaternion([tgt_q.x, tgt_q.y, tgt_q.z, tgt_q.w]))
        cal_e = np.degrees(tft.euler_from_quaternion([cal_q.x, cal_q.y, cal_q.z, cal_q.w]))
        diff = np.abs(tgt_e - cal_e)
        rospy.logdebug(f"姿态偏差: {diff}")
        return np.all(diff <= self.orientation_tol)

    # ========== 控制输出 ==========
    def publish_joint_traj(self, joint_state):
        traj = JointTrajectory()
        traj.joint_names = joint_state.name
        pt = JointTrajectoryPoint()
        pt.positions = joint_state.position
        pt.velocities = [0.0] * len(pt.positions)
        pt.time_from_start = rospy.Duration(1.0 / self.loop_rate)
        traj.points = [pt]
        self.joint_pub.publish(traj)

    # ========== 辅助函数 ==========
    @staticmethod
    def quat_to_rot(q):
        return tft.quaternion_matrix([q.x, q.y, q.z, q.w])[:3, :3]

    @staticmethod
    def print_pose(pose, prefix="", throttle=None):
        p = pose.pose.position
        q = pose.pose.orientation
        roll, pitch, yaw = np.degrees(tft.euler_from_quaternion([q.x, q.y, q.z, q.w]))  # 注意顺序
        msg = f"{prefix}: x={p.x:.3f}, y={p.y:.3f}, z={p.z:.3f}, roll={roll:.1f}°, pitch={pitch:.1f}°, yaw={yaw:.1f}°"
        if throttle:
            rospy.loginfo_throttle(throttle, msg)
        else:
            rospy.loginfo(msg)


if __name__ == "__main__":
    try:
        TeleopPose()
    except rospy.ROSInterruptException:
        pass
    finally:
        roscpp_shutdown()
